FM: Generative Models via Flow Matching
Flow matching is an amazingly simple algorithm for generative models. It has been widely adopted in image generation models.
Table of Contents
1. Motivation
In Diffusion models, we start from a sampled Gaussian noise, and then take multiple steps to gradually denoise to obtain a clean image.
How about directly transforming a noise into a clean image? Can’t we just interpolate linearly between noise and image? The answer is yes, and that’s exactly the idea of flow matching. Flow matching (FM hereafter) directly learns the path from noise to data. More specifically, FM learns a velocity field \(v_{t}(\mathbf{x})\), the velocity (including speed and direction) of particle \(\mathbf{x}\) (i.e., the image) at timestamp \(t \in [0,1]\) and at position \(\mathbf{x}\). The velocity represents how the image is “denoised”. Here, \(t=0\) is the complete noise and \(t=1\) is the clean image. The inference (image generation process) is simple: we first sample a noise at \(t=0\), then integrate along \(v_{t}(\mathbf{x})\) over \(t\), and we’ll obtain the image.
2. Continuous Normalizing Flow
FM is built on continuous normalizing flow (CNF)
Ideally, if we have already obtained the true velocity field \(u_{t}(\mathbf{x})\), then we can simply do a regression:
\[ \mathcal{L}_{FM}(\theta) = \mathop{\mathbb{E}}\limits_{\begin{subarray}{c} t \sim\mathcal{U}[0,1]\\ x \sim p_{t}(x) \end{subarray}} \left[ \left\| v_{\theta}(x,t) - u_{t}(x) \right\|^{2} \right] \]
The issue is that, we don’t know \(u_{t}(x)\). It turns out that, we have an alternative: conditioned flow matching.
3. Condition Flow Matching (CFM)
Since we don’t know \(u_{t}(x)\), we can optimize using the conditioned version of velocity field \(u_{t}(x|x_{1})\) conditioned by ground truth \(x_{1}\). We define the CFM loss
\[ \mathcal{L}_{CFM}(\theta) = \mathop{\mathbb{E}}\limits_{\begin{subarray}{c} t\sim\mathcal{U}[0,1]\\ x_{1}\sim p_{data}(x)\\ x\sim p_{t}(x|x_{1}) \end{subarray}} \left[ \left\| v_{\theta}(x,t) - u_{t}(x|x_{1}) \right\|^{2} \right] \]
Indeed, we can mathematically prove that \(\nabla_{\theta} \mathcal{L}_{CFM} = \nabla_{\theta} \mathcal{L}_{FM}\). Therefore, at last, \(v_{\theta}\) learnt with CFM loss will converge to the correct \(u_{t}(x)\) one learnt with FM loss.
We’ll prove \(\nabla_{\theta} \mathcal{L}_{CFM}(\theta) = \nabla_{\theta} \mathcal{L}_{FM}(\theta)\).
3.1. Using Linear Interpolation to Simplify CFM
The simplest path from noise to clear image is linear path, therefore
\[ u_{t}(x|x_{1}) = x_{1}-x_{0}=\frac{x_{1}-x_{t}}{1-t}, \quad t \in [0,1] \]
Therefore, the CFM loss now becomes
\[ \mathcal{L}_{CFM}(\theta) = \mathbb{E}_{t,x_{0},x_{1}}\left[ \left\| v_{\theta}(x_{t}, t) - (x_{1} - x_{0}) \right\|^{2} \right] \]
4. Flow Matching Algorithm
So now, we can describe the training and inference algorithm with flow matching.
4.1. Training
- Sample a ground truth \(x_{1}\) from dataset
- Sample a noise \(x_{0}\) from standard Gaussian
- Sample a timestamp \(t\) from \(\mathcal{U}[0,1]\)
- Construct the linear interpolation \(x_{t} = (1-t)x_{0} + tx_{1}\)
- Compute \(u_{t}(x|xe_{1})=x_{1}-x_{0}\)
- Compute loss \(\| v_{\theta}(x_{t}, t) - u_{t}(x|x_{1}) \|^{2}\) and update the parameters.
4.2. Inference
- Sample an initial noise \(x_{0}\)
- Use any ODE solver to integrate from \(t=0\) to \(t=1\)
- At each step, compute \(v_{\theta}(x_{t},t)\) to compute velocity and advance state
- The final \(x_{1}\) is the generated data